#include "MultigridSolver.h"
#include "Grid.h"
#include "BoundaryConditions.h"
#include <iostream>
#include <fstream>
#include <nlohmann/json.hpp>
#include <chrono>
#include <iomanip>
#include <vector>
#include <cmath>
#include <string>
#include <filesystem>

using json = nlohmann::json;
using namespace Eigen;

void run_test_case(const std::string& config_path, 
    const std::string& case_name) {
    std::cout << "\n=== Running Test Case: " << case_name << " ===\n";

    // 加载配置
    std::ifstream config_file(config_path);
    json config = json::parse(config_file);

    // 网格尺寸测试
    const std::vector<int> grid_sizes = {32, 64, 128, 256};
    std::vector<double> l2_errors;      // 存储L2误差
    std::vector<double> max_errors;     // 存储最大范数误差
    std::vector<double> l2_rates;       // L2收敛率
    std::vector<double> max_rates;      // 最大范数收敛率

    for (size_t i = 0; i < grid_sizes.size(); ++i) {
        const int n = grid_sizes[i];
        config["grid"]["size"] = n;

        // 写入临时配置
        std::ofstream temp_config("temp_config.json");
        temp_config << config.dump(4);
        temp_config.close();

        // 初始化求解器
        Grid grid("temp_config.json");
        BoundaryConditions bc("temp_config.json");
        MultigridSolver solver(grid, bc, "temp_config.json");

        // 求解并计时
        auto start = std::chrono::steady_clock::now();
        VectorXd solution = solver.solve();
        auto end = std::chrono::steady_clock::now();
        std::chrono::duration<double> elapsed = end - start;

        // 计算L2误差
        double l2_error = solver.compute_error();
        l2_errors.push_back(l2_error);

        // 计算最大范数误差
        VectorXd exact_solution = solver.getExact_solution();
        VectorXd error_vector = exact_solution - solution;
        double max_error = error_vector.lpNorm<Infinity>();
        max_errors.push_back(max_error);

        // 输出结果
        std::cout << "\nGrid Size: " << std::setw(4) << n 
            << " | L2 Error: " << std::scientific << std::setprecision(3) << l2_error
            << " | Max Error: " << std::scientific << std::setprecision(3) << max_error
            << " | Time: " << std::fixed << std::setprecision(4) << elapsed.count() << "s\n";

        // 计算收敛率（至少第二个网格开始）
        if (i > 0) {
            double l2_rate = log2(l2_errors[i-1] / l2_errors[i]);
            l2_rates.push_back(l2_rate);

            double max_rate = log2(max_errors[i-1] / max_errors[i]);
            max_rates.push_back(max_rate);

            std::cout << "  L2 Convergence Rate: " << std::fixed << std::setprecision(2) << l2_rate << "\n";
            std::cout << "  Max Norm Convergence Rate: " << std::fixed << std::setprecision(2) << max_rate << "\n";
        }

        // 输出解到 CSV 文件
        std::string output_dir = "output";
        std::string filename = output_dir + "/solution_" + case_name + "_" + std::to_string(n) + ".csv";
        std::ofstream output_file(filename);
        if (output_file.is_open()) {
            for (int j = 0; j < solution.size(); ++j) {
                output_file << solution(j) << "\n";
            }
            output_file.close();
        }
    }

    // 输出平均收敛率
    if (!l2_rates.empty()) {
        double avg_l2_rate = std::accumulate(l2_rates.begin(), l2_rates.end(), 0.0) / l2_rates.size();
        double avg_max_rate = std::accumulate(max_rates.begin(), max_rates.end(), 0.0) / max_rates.size();

        std::cout << "\nAverage Convergence Rates:\n";
        std::cout << "  L2: " << std::fixed << std::setprecision(2) << avg_l2_rate << "\n";
        std::cout << "  Max Norm: " << std::fixed << std::setprecision(2) << avg_max_rate << "\n";
    }
}

int main(int argc, char* argv[]) {
    try {
        const std::string config_file = (argc > 1) ? argv[1] : "input/config.json";

        std::string case_name;
        if (argc > 1) {
            // 将 argv[1] 转换为 std::string
            std::string config_path = argv[1];

            // 提取文件名作为 case_name
            size_t last_slash_pos = config_path.find_last_of("/\\");
            if (last_slash_pos != std::string::npos) {
                case_name = config_path.substr(last_slash_pos + 1);
            } else {
                case_name = config_path;
            }

            // 去掉文件扩展名
            size_t last_dot_pos = case_name.find_last_of(".");
            if (last_dot_pos != std::string::npos) {
                case_name = case_name.substr(0, last_dot_pos);
            }
        } else {
            // 如果没有额外的参数，使用默认值 "config"
            case_name = "config";
        }

        // 打印 case_name
        std::cout << "Using case name: " << case_name << std::endl;

        run_test_case(config_file, case_name);

    } catch (const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}